{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "9ab64209",
   "metadata": {},
   "outputs": [],
   "source": [
    "from openvino.runtime import Core\n",
    "import openvino.runtime as ov\n",
    "import cv2 as cv\n",
    "import numpy as np\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "fbe93bbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import paddle.fluid as fluid"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6239d3e",
   "metadata": {},
   "source": [
    "## OpenVINO 模型推理器（class）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "be4b7342",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Predictor:\n",
    "    \"\"\"\n",
    "    OpenVINO 模型推理器\n",
    "    \"\"\"\n",
    "    def __init__(self, model_path):\n",
    "        ie_core = Core()\n",
    "        model = ie_core.read_model(model=model_path)\n",
    "        self.compiled_model = ie_core.compile_model(model=model, device_name=\"CPU\")\n",
    "    def get_inputs_name(self, num):\n",
    "        return self.compiled_model.input(num)\n",
    "    \n",
    "    def get_outputs_name(self, num):\n",
    "        return self.compiled_model.output(num)\n",
    "    \n",
    "    def predict(self, input_data):\n",
    "        return self.compiled_model([input_data])\n",
    "    \n",
    "    def get_request(self):\n",
    "        return self.compiled_model.create_infer_request()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21b888e3",
   "metadata": {},
   "source": [
    "## 图像预处理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "b441c64e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_image(input_image, size):\n",
    "    \"\"\"输入图片与处理方法，按照PP-Yoloe模型要求预处理图片数据\n",
    "\n",
    "    Args:\n",
    "        input_image (uint8): 输入图片矩阵\n",
    "        size (int): 模型输入大小\n",
    "\n",
    "    Returns:\n",
    "        float32: 返回处理后的图片矩阵数据\n",
    "    \"\"\"\n",
    "    max_len = max(input_image.shape)\n",
    "    img = np.zeros([max_len,max_len,3],np.uint8)\n",
    "    img[0:input_image.shape[0],0:input_image.shape[1]] = input_image # 将图片放到正方形背景中\n",
    "    img = cv.cvtColor(img,cv.COLOR_BGR2RGB)  # BGR转RGB\n",
    "    img = cv.resize(img, (size, size), cv.INTER_NEAREST) # 缩放图片\n",
    "    img = np.transpose(img,[2, 0, 1]) # 转换格式\n",
    "    img = img / 255.0 # 归一化\n",
    "    img = np.expand_dims(img,0) # 增加维度\n",
    "    return img.astype(np.float32)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50d7ff5e",
   "metadata": {},
   "source": [
    "## 图像后处理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "660d45cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_result(box_results, conf_results):\n",
    "    \"\"\"按照PP-Yolove模型输出要求，处理数据，非极大值抑制，提取预测结果\n",
    "\n",
    "    Args:\n",
    "        box_results (float32): 预测框预测结果\n",
    "        conf_results (float32): 置信度预测结果\n",
    "    Returns:\n",
    "        float: 预测框\n",
    "        float: 分数\n",
    "        int: 类别\n",
    "    \"\"\"\n",
    "    conf_results = np.transpose(conf_results,[0, 2, 1]) # 转置\n",
    "    # 设置输出形状\n",
    "    box_results =box_results.reshape(8400,4) \n",
    "    conf_results = conf_results.reshape(8400,2)\n",
    "    scores = []\n",
    "    classes = []\n",
    "    boxes = []\n",
    "    for i in range(8400):\n",
    "        conf = conf_results[i,:] # 预测分数\n",
    "        score = np.max(conf) # 获取类别\n",
    "        # 筛选较小的预测类别\n",
    "        if score > 0.5:\n",
    "            classes.append(np.argmax(conf)) \n",
    "            scores.append(score) \n",
    "            boxes.append(box_results[i,:])\n",
    "    scores = np.array(scores)\n",
    "    boxes = np.array(boxes)\n",
    "    \n",
    "    result_box = []\n",
    "    result_score = []\n",
    "    result_class = []\n",
    "    # 非极大值抑制筛选重复的预测结果\n",
    "    if len(boxes) != 0:\n",
    "        # 非极大值抑制结果\n",
    "        indexs = tf.image.non_max_suppression(boxes,scores,len(scores),0.25,0.35)\n",
    "        for i, index in enumerate(indexs):\n",
    "            result_score.append(scores[index])\n",
    "            result_box.append(boxes[index,:])\n",
    "            result_class.append(classes[index])\n",
    "    # 返回结果\n",
    "    return np.array(result_box),np.array(result_score),np.array(result_class)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54d90f67",
   "metadata": {},
   "source": [
    "## 画出预测框"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "212c1bde",
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_box(image, boxes, scores, classes, labels):\n",
    "    \"\"\"将预测结果绘制到图像上\n",
    "\n",
    "    Args:\n",
    "        image (uint8): 原图片\n",
    "        boxes (float32): 预测框\n",
    "        scores (float32): 分数\n",
    "        classes (int): 类别\n",
    "        lables (str): 标签\n",
    "\n",
    "    Returns:\n",
    "        uint8: 标注好的图片\n",
    "    \"\"\"\n",
    "    colors = [(0, 0, 255), (0, 255, 0)]\n",
    "    scale = max(image.shape) / 640.0 # 缩放比例\n",
    "    if len(classes) != 0:\n",
    "        for i in range(len(classes)):\n",
    "            box = boxes[i,:]\n",
    "            x1 = int(box[0] * scale)\n",
    "            y1 = int(box[1] * scale)\n",
    "            x2 = int(box[2] * scale)\n",
    "            y2 = int(box[3] * scale)\n",
    "            label = labels[classes[i]]\n",
    "            score = scores[i]\n",
    "            cv.rectangle(image, (x1, y1), (x2, y2), colors[classes[i]], 2, cv.LINE_8)\n",
    "            cv.putText(image,label+\":\"+str(score),(x1,y1-10),cv.FONT_HERSHEY_SIMPLEX, 0.55, colors[classes[i]], 2)\n",
    "        \n",
    "    return image\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed16c8ef",
   "metadata": {},
   "source": [
    "## 读取标签"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "2eb56f59",
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_label(label_path):\n",
    "    with open(label_path, 'r') as f:\n",
    "        labels = f.read().split()\n",
    "    return labels"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26bb2326",
   "metadata": {},
   "source": [
    "## 同步推理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "d2898cc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "label_path = \"labels.txt\"\n",
    "yoloe_model_path = \"ppyoloe_crn_s_80.xml\"\n",
    "predictor = Predictor(model_path = yoloe_model_path)\n",
    "boxes_name = predictor.get_outputs_name(0)\n",
    "conf_name = predictor.get_outputs_name(1)\n",
    "labels = read_label(label_path=label_path)\n",
    "cap = cv.VideoCapture(0)\n",
    "while cap.isOpened():\n",
    "    ret, frame = cap.read()\n",
    "    frame = cv.flip(frame, 180)\n",
    "    cv.namedWindow(\"MaskDetection\", 0)  # 0可调大小，注意：窗口名必须imshow里面的一窗口名一直\n",
    "    cv.resizeWindow(\"MaskDetection\", 640, 480)    # 设置长和宽\n",
    "    input_frame = process_image(frame, 640)\n",
    "    results = predictor.predict(input_data=input_frame)\n",
    "    boxes, scores, classes = process_result(box_results=results[boxes_name], conf_results=results[conf_name])\n",
    "    result_frame = draw_box(image=frame, boxes=boxes, scores=scores, classes=classes, labels=labels)\n",
    "    cv.imshow('MaskDetection', result_frame)\n",
    "    key = cv.waitKey(1)\n",
    "    if key == 27: #esc退出\n",
    "        break\n",
    "cap.release()\n",
    "cv.destroyAllWindows()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94cb6ccb",
   "metadata": {},
   "source": [
    "## 异步推理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "a7b35607",
   "metadata": {},
   "outputs": [],
   "source": [
    "label_path = \"labels.txt\"\n",
    "yoloe_model_path = \"ppyoloe_crn_s_80.xml\"\n",
    "predictor = Predictor(model_path = yoloe_model_path)\n",
    "input_layer = predictor.get_inputs_name(0)\n",
    "labels = read_label(label_path=label_path)\n",
    "cap = cv.VideoCapture(0)\n",
    "curr_request = predictor.get_request()\n",
    "next_request = predictor.get_request()\n",
    "ret, frame = cap.read()\n",
    "curr_frame = process_image(frame, 640)\n",
    "curr_request.set_tensor(input_layer, ov.Tensor(curr_frame))\n",
    "curr_request.start_async()\n",
    "while cap.isOpened():\n",
    "    ret, next_frame = cap.read()\n",
    "    next_frame = cv.flip(next_frame, 180)\n",
    "    cv.namedWindow(\"MaskDetection\", 0)  # 0可调大小，注意：窗口名必须imshow里面的一窗口名一直\n",
    "    cv.resizeWindow(\"MaskDetection\", 640, 480)    # 设置长和宽\n",
    "    in_frame = process_image(next_frame, 640)\n",
    "    next_request.set_tensor(input_layer, ov.Tensor(in_frame))\n",
    "    next_request.start_async()\n",
    "    if curr_request.wait_for(-1) == 1:\n",
    "        boxes_name = curr_request.get_output_tensor(0).data\n",
    "        conf_name = curr_request.get_output_tensor(1).data\n",
    "        boxes, scores, classes = process_result(box_results=boxes_name, conf_results=conf_name)\n",
    "        frame = draw_box(image=frame, boxes=boxes, scores=scores, classes=classes, labels=labels)\n",
    "        cv.imshow('MaskDetection', frame)\n",
    "    frame = next_frame\n",
    "    curr_request, next_request = next_request, curr_request\n",
    "    key = cv.waitKey(1)\n",
    "    if key == 27: #esc退出\n",
    "        break\n",
    "cap.release()\n",
    "cv.destroyAllWindows()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6547f88",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
